import cv2 as cv
import numpy as np
import pytesseract
from StreamProcessor import VideoProcessor


class ChiOcr:
    # 清洗识别字符串
    @staticmethod
    def drop(reg_str: str):
        reg_list = list(reg_str.strip())
        while ' ' in reg_list:
            reg_list.remove(' ')
        return ''.join(reg_list)

    # 根据轮廓对给定的图像进行仿射变换矫正
    @staticmethod
    def affine(img: np.ndarray, cnt):
        rect = cv.minAreaRect(cnt)
        (x, y), (w, h), (angle) = rect
        y, x = img.shape[:2]
        M = cv.getRotationMatrix2D((x // 2, y // 2), angle - 90, 1.0)
        cos = np.abs(M[0, 0])
        sin = np.abs(M[0, 1])
        nW = int((w * sin) + (h * cos))
        nH = int(w)
        M[0, 2] += (nW // 2) - x // 2
        M[1, 2] += (nH // 2) - y // 2

        return cv.warpAffine(img, M, (nW, nH), flags=cv.INTER_CUBIC, borderMode=cv.BORDER_REPLICATE)

    # 文字检测并仿射变换返回矫正文字
    def text_detect(self, img: np.ndarray):
        equ = cv.equalizeHist(img)  # 直方图均衡化增强图像
        thresh = cv.adaptiveThreshold(equ, 255, cv.ADAPTIVE_THRESH_GAUSSIAN_C, cv.THRESH_BINARY, 3, 14)  # 自适应阈值化
        sobelY = cv.Sobel(thresh, cv.CV_8UC1, 0, 1, ksize=9)  # sobel算子边缘检测
        closing = cv.morphologyEx(sobelY, cv.MORPH_DILATE, np.ones((6, 6)))  # 文字感兴趣区域ROI检测
        contours = cv.findContours(closing, cv.RETR_EXTERNAL, cv.CHAIN_APPROX_SIMPLE)[0]  # ROI轮廓提取

        ROIs = []
        for cnt in contours:
            x, y, w, h = cv.boundingRect(cnt)
            if w < 50 or h < 50:  # ROI筛选
                continue
            ROI = img[y:y + h, x:x + w]
            ROI = self.affine(ROI, cnt)  # 仿射变换
            ROIs.append(ROI)

        return ROIs

    # tesseract识别矫正后文字
    def ocr(self, processor: VideoProcessor):
        while True:
            try:
                if len(processor.img[0]) != 0:
                    img = cv.cvtColor(processor.img[0], cv.COLOR_BGR2GRAY)
                    ROIs = self.text_detect(img)
                    for ROI in ROIs:
                        thresh = cv.threshold(ROI, 0, 255, cv.THRESH_BINARY + cv.THRESH_OTSU)[1]
                        recg_str = pytesseract.image_to_string(thresh, "chi_sim")
                        if len(self.drop(recg_str)) != 0:
                            print(self.drop(recg_str))
            except IndexError:
                pass
